import os
import pdb
import numpy as np
import csv

os.chdir('/home/wjh/Meta_R-CNN/txt/')

classes = ['background', 'aeroplane', 'bicycle', 'boat', 'bottle', 'car', 'cat', 'chair', 'diningtable', 'dog',
                         'horse', 'person', 'pottedplant', 'sheep', 'train', 'tvmonitor', 'bird', 'bus', 'cow',
                         'motorbike', 'sofa']
                         
txt_list1 = ['gt_cls.txt', 'sim_cls.txt']
txt_list2 = ['gt_cls_all.txt', 'pre_cls.txt', 're_pre_cls.txt']

def read_txt(txt_file):
    data_list = []
    with open(txt_file, 'r') as f:
        for line in f.readlines():
            line = line.strip('\n')
            data_list.append(int(float(line)))
    
    return data_list

cls_list1 = []
cls_list2 = []

print('Reading txt file ...')
'''
for txt_file in txt_list1:
    cls_list1.append(read_txt(txt_file))
'''
for txt_file in txt_list2:
    cls_list2.append(read_txt(txt_file))

print('Finished read.')

cls_matrix = [[[] for _ in range(len(classes))] for _ in range(len(classes))]


gt_cls = np.array(cls_list1[0])
sim_cls = np.array(cls_list1[1])

print('Evaluating similarity matrix ... ')
csvfile = open('similarity_matrix.csv', 'w')
writer = csv.writer(csvfile)
writer.writerow([' '] + classes)
for i in range(len(classes)):
    for j in range(len(classes)):
        cls_matrix[i][j] = sum(sim_cls[np.where(gt_cls == i)] == j)
    writer.writerow([classes[i]] + cls_matrix[i])
csvfile.close()
print('Done.')


gt_cls_all = np.array(cls_list2[0])
pre_cls = np.array(cls_list2[1])
re_pre_cls = np.array(cls_list2[2])

print('Evaluating reweight matrix ... ')
csvfile = open('reweight_matrix.csv', 'w')
writer = csv.writer(csvfile)
writer.writerow([' '] + classes)
for i in range(len(classes)):
    for j in range(len(classes)):
        cls_matrix[i][j] = np.sum(re_pre_cls[np.where(pre_cls == i)] == j)
    writer.writerow([classes[i]] + cls_matrix[i])
csvfile.close()
print('Done.')

# pdb.set_trace()
    